import torch
import numpy as np


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def bresenham(x1, y1, x2, y2):
    """
    Bresenham algorithm in all four quadrants
    :param x1, y1: coordinate of the lower left point
    :param x2, y2: coordinate of the upper right point
    :return: a list of the points on the line segment from (x1, y1) to (x2, y2)
                        (x2, y2)
    -----------------/--------------------------
    |              /                            |
    -------------/------------------------------
    |          /                                |
    _________/___________________________________
    |      /                                    |
    _____/______________________________________
    |  /                                        |
    _/__________________________________________
    (x1,y1)
    In the condition of slope ~ (0, 1], just compute the decision parameter P,
    initial state P1 = 2dy-dx, A = 2dy, B = 2dy - 2dx. Use these values to update
    the (xk, yk) is enough, but for all the slope, we need two tricks:
    1. |m| <= 1, no interchange
    2. |m| > 1, interchange x and y
    3. sign(m) > 0, increase x(y) every time depend on the interchange
    4. sign(m) <= 0, decrease x(y) every time depend on the interchange
    """

    # compute P, P1, and which condition is the line segment is in
    x = x1
    y = y1
    dx = np.abs(x2 - x1)
    dy = np.abs(y2 - y1)
    sx = np.sign(x2 - x1)
    sy = np.sign(y2 - y1)
    if dy - dx > 0:
        dx, dy = dy, dx
        interchange = True
    else:
        interchange = False


    E = 2*dy - dx
    A = 2*dy
    B = 2*(dy - dx)

    # compute the point list, defined by x_bres and y_bres
    x_bres, y_bres = [], []
    x_bres.append(x)
    y_bres.append(y)
    for _ in range(1, int(dx)):
        if E > 0:
            y += sy
            x += sx
            E += B
        else:
            if interchange:
                y += sy
            else:
                x += sx
            E += A

        x_bres.append(x)
        y_bres.append(y)
    return zip(x_bres, y_bres)


def bresenham_torch(start, end, max_iter = 5):
    """
    The fake bresenham algorithm implemented in pytorch,
    though it is fake, but it is batched, and floats running on gpus
    does not matter
    """
    return _bresenhamlines(start, end, max_iter).reshape(-1, start.shape[-1])

def _bresenhamline_nslope(slope):
    """
    Normalize slope for Bresenham's line algorithm.

    >>> s = np.array([[-2, -2, -2, 0]])
    >>> _bresenhamline_nslope(s)
    array([[-1., -1., -1.,  0.]])

    >>> s = np.array([[0, 0, 0, 0]])
    >>> _bresenhamline_nslope(s)
    array([[ 0.,  0.,  0.,  0.]])

    >>> s = np.array([[0, 0, 9, 0]])
    >>> _bresenhamline_nslope(s)
    array([[ 0.,  0.,  1.,  0.]])
    """
    scale = torch.amax(torch.abs(slope), dim=1).reshape(-1, 1)
    zeroslope = (scale == 0).all(1)
    scale[zeroslope] = torch.ones(1).to(device)
    normalizedslope = slope / scale
    normalizedslope[zeroslope] = torch.zeros(slope[0].shape).to(device)
    return normalizedslope

def _bresenhamlines(start, end, max_iter):
    """
    Returns npts lines of length max_iter each. (npts x max_iter x dimension) 

    >>> s = np.array([[3, 1, 9, 0],[0, 0, 3, 0]])
    >>> _bresenhamlines(s, np.zeros(s.shape[1]), max_iter=-1)
    array([[[ 3,  1,  8,  0],
            [ 2,  1,  7,  0],
            [ 2,  1,  6,  0],
            [ 2,  1,  5,  0],
            [ 1,  0,  4,  0],
            [ 1,  0,  3,  0],
            [ 1,  0,  2,  0],
            [ 0,  0,  1,  0],
            [ 0,  0,  0,  0]],
    <BLANKLINE>
           [[ 0,  0,  2,  0],
            [ 0,  0,  1,  0],
            [ 0,  0,  0,  0],
            [ 0,  0, -1,  0],
            [ 0,  0, -2,  0],
            [ 0,  0, -3,  0],
            [ 0,  0, -4,  0],
            [ 0,  0, -5,  0],
            [ 0,  0, -6,  0]]])
    """
    if max_iter == -1:
        max_iter = torch.amax(torch.amax(torch.abs(end - start), dim=1))
    npts, dim = start.shape
    nslope = _bresenhamline_nslope(end - start)

    # steps to iterate on
    stepseq = torch.arange(1, max_iter + 1).to(device)
    stepmat = stepseq.repeat(dim, 1) #np.tile(stepseq, (dim, 1)).T
    stepmat = stepmat.T

    # some hacks for broadcasting properly
    bline = start[:, None, :] + nslope[:, None, :] * stepmat

    # Approximate to nearest int
    bline_points = torch.round(bline).to(start.dtype)
    
    return bline_points